import argparse
import csv
import json
import logging
import random
import sys
from pathlib import Path
from typing import Dict, List, Any, Optional, Tuple

# 导入第三方库，如果未安装，需要 pip install pyyaml
try:
    import yaml
except ImportError:
    print("Error: PyYAML is not installed. Please install it using 'pip install pyyaml'", file=sys.stderr)
    sys.exit(1)

# 配置日志
# 确保日志配置在模块级别，并且只配置一次
if not logging.getLogger(__name__).handlers:
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class InputGeneratorError(Exception):
    """
    自定义异常类，用于表示输入文件生成过程中发生的错误。

    Attributes:
        message (str): 错误描述信息。
        error_code (int): 错误代码，便于分类处理。
    """
    def __init__(self, message: str, error_code: int = 1):
        self.message = message
        self.error_code = error_code
        super().__init__(self.message)

    def __str__(self) -> str:
        return f"[{self.error_code}] {self.message}"


def read_fasta(file_path: str) -> str:
    """
    读取FASTA文件并返回找到的第一个序列。

    Args:
        file_path (str): FASTA文件路径。

    Returns:
        str: 提取到的序列字符串。

    Raises:
        InputGeneratorError: 当FASTA文件未找到、无法读取或文件中未找到序列时抛出。
    """
    sequence_lines: List[str] = []
    try:
        with open(file_path, 'r') as f:
            for line in f:
                line = line.strip()
                if not line:  # 跳过空行
                    continue
                if line.startswith('>'):
                    if sequence_lines:  # 如果已经获取到序列，遇到下一个标题时停止
                        break
                    continue
                sequence_lines.append(line)
    except FileNotFoundError:
        raise InputGeneratorError(f"FASTA文件未找到: {file_path}", error_code=101)
    except Exception as e:
        raise InputGeneratorError(f"读取FASTA文件失败: {file_path}, 错误: {str(e)}", error_code=102)

    if not sequence_lines:
        raise InputGeneratorError(f"在FASTA文件中未找到序列: {file_path}", error_code=103)
    return ''.join(sequence_lines)


def read_compound_csv(file_path: str) -> List[Dict[str, str]]:
    """
    读取包含化合物信息的CSV文件。

    Args:
        file_path (str): CSV文件路径。

    Returns:
        List[Dict[str, str]]: 化合物信息列表，每个元素为包含'id'和'smiles'的字典。

    Raises:
        InputGeneratorError: 当CSV文件未找到、无法读取、缺少必要列或未找到有效化合物数据时抛出。
    """
    compounds: List[Dict[str, str]] = []
    try:
        with open(file_path, 'r', newline='') as f:
            reader = csv.DictReader(f)
            # 检查是否存在必要的列
            required_columns = {'id', 'smiles'}
            if not reader.fieldnames or not required_columns.issubset(reader.fieldnames):
                missing = required_columns - set(reader.fieldnames or [])
                raise InputGeneratorError(f"CSV文件缺少必要的列: {', '.join(missing)}", error_code=104)

            for row_num, row in enumerate(reader, 2):  # 行号从2开始(标题行是1)
                compound_id = row.get('id', '').strip()
                smiles = row.get('smiles', '').strip()
                if not compound_id or not smiles:
                    logger.warning(f"CSV文件 '{file_path}' 第{row_num}行数据不完整（id或smiles为空），已跳过。")
                    continue
                compounds.append({'id': compound_id, 'smiles': smiles})

            if not compounds:
                raise InputGeneratorError(f"CSV文件 '{file_path}' 中未找到有效的化合物数据。", error_code=105)
            return compounds
    except FileNotFoundError:
        raise InputGeneratorError(f"CSV文件未找到: {file_path}", error_code=106)
    except Exception as e:
        raise InputGeneratorError(f"读取CSV文件失败: {file_path}, 错误: {str(e)}", error_code=107)


def generate_random_seed(seed_path: str, min_value: int = 101, max_value: int = 999999) -> int:
    """
    生成随机种子并写入文件。

    Args:
        seed_path (str): 种子文件保存路径。
        min_value (int): 随机种子最小值。
        max_value (int): 随机种子最大值。

    Returns:
        int: 生成的随机种子值。

    Raises:
        InputGeneratorError: 当写入种子文件失败时抛出。
    """
    seed = random.randint(min_value, max_value)
    try:
        Path(seed_path).parent.mkdir(parents=True, exist_ok=True)
        with open(seed_path, 'w') as f_seed:
            f_seed.write(str(seed))
        logger.info(f"随机种子已生成并写入: {seed_path} (种子值: {seed})")
        return seed
    except Exception as e:
        raise InputGeneratorError(f"写入种子文件失败: {seed_path}, 错误: {str(e)}", error_code=108)


class BaseCommand:
    """
    命令基类，提供通用功能和接口定义。
    所有具体命令（如Boltz, Protenix）都应继承此类。
    """
    def __init__(self, args: argparse.Namespace):
        self.args = args

    @staticmethod
    def add_arguments(parser: argparse.ArgumentParser) -> None:
        """
        添加特定命令的参数。子类必须实现此方法。

        Args:
            parser (argparse.ArgumentParser): 参数解析器对象。
        """
        raise NotImplementedError("子类必须实现 add_arguments 方法来定义其命令行参数。")

    def run(self) -> None:
        """
        执行命令的主方法。子类必须实现此方法。
        """
        raise NotImplementedError("子类必须实现 run 方法来执行其特定逻辑。")

    def _parse_sequence_arg(self, sequence_arg_str: Optional[str], seq_type_key: str, id_prefix: str,
                           start_id: int, add_id: bool = True) -> List[Dict[str, Any]]:
        """
        解析命令行中的序列参数字符串（例如 "path:count,path2:count2"）。

        Args:
            sequence_arg_str (Optional[str]): 命令行参数字符串，例如 "file1.fasta:2,file2.fasta:3"。
            seq_type_key (str): 序列类型在输出字典中的键名 (e.g., 'protein', 'dna')。
            id_prefix (str): 序列ID的前缀 (e.g., 'P', 'D', 'R')。
            start_id (int): 用于生成唯一序列ID的起始数字 (仅当 add_id 为 True 时使用)。
            add_id (bool): 是否在序列字典中添加 'id' 字段。默认为 True。

        Returns:
            List[Dict[str, Any]]: 包含解析后的序列信息的列表。

        Raises:
            InputGeneratorError: 当序列格式错误或FASTA文件无法读取时抛出。
        """
        sequences: List[Dict[str, Any]] = []
        if sequence_arg_str:
            for item in sequence_arg_str.split(','):
                try:
                    path, count_str = item.strip().split(':')
                    count = int(count_str.strip())
                    if count <= 0:
                        raise ValueError("数量必须是正整数")
                    
                    seq = read_fasta(path.strip())
                    
                    item_data: Dict[str, Any] = {
                        'sequence': seq,
                        'count': count
                    }
                    if add_id:
                        item_data['id'] = f"{id_prefix}{start_id}"
                        start_id += 1 # 只有在添加ID时才递增ID计数器
                    
                    sequences.append({seq_type_key: item_data})
                except ValueError as e:
                    raise InputGeneratorError(
                        f"序列参数格式错误或数量无效 '{item}'，应为'路径:数量'且数量为正整数: {str(e)}",
                        error_code=202
                    ) from e
                except InputGeneratorError as e:
                    # Re-raise InputGeneratorError from read_fasta
                    raise InputGeneratorError(f"处理序列文件 '{path}' 失败: {e.message}", e.error_code) from e
        return sequences

    def _collect_sequences(self, add_id: bool = True) -> Tuple[List[Dict[str, Any]], int]:
        """
        收集并解析所有序列信息（蛋白质、DNA、RNA）。

        Args:
            add_id (bool): 是否在序列字典中添加 'id' 字段。默认为 True。

        Returns:
            Tuple[List[Dict[str, Any]], int]: 包含各类序列信息的列表，以及下一个可用的序列ID。
                                               如果 add_id 为 False，则返回的 int 值无意义。

        Raises:
            InputGeneratorError: 当未提供任何序列时抛出。
        """
        all_sequences: List[Dict[str, Any]] = []
        current_id = 1  # 用于生成唯一序列ID的计数器
        
        # 确定蛋白质序列的键名
        # Protenix 命令需要使用 'proteinChain'，其他命令使用 'protein'
        protein_key = 'proteinChain' if isinstance(self, ProtenixCommand) else 'protein'

        protein_seqs = self._parse_sequence_arg(self.args.protein, protein_key, 'P', current_id, add_id)
        if add_id:
            current_id += len(protein_seqs)

        dna_seqs = self._parse_sequence_arg(self.args.dna, 'dna', 'D', current_id, add_id)
        if add_id:
            current_id += len(dna_seqs)

        rna_seqs = self._parse_sequence_arg(self.args.rna, 'rna', 'R', current_id, add_id)
        if add_id:
            current_id += len(rna_seqs)

        all_sequences.extend(protein_seqs)
        all_sequences.extend(dna_seqs)
        all_sequences.extend(rna_seqs)

        if not all_sequences:
            raise InputGeneratorError("必须提供至少一个序列文件（--protein, --dna 或 --rna）", error_code=201)
        return all_sequences, current_id

    def _parse_ccd_ligands(self, add_id: bool = True) -> List[Dict[str, Any]]:
        """
        解析CCD配体参数。

        Args:
            add_id (bool): 是否在CCD配体字典中添加 'id' 字段。默认为 True。

        Returns:
            List[Dict[str, Any]]: CCD配体信息列表。

        Raises:
            InputGeneratorError: 当CCD配体格式错误时抛出。
        """
        ccd_ligands: List[Dict[str, Any]] = []
        if self.args.ccd:
            try:
                ligand_parts = self.args.ccd.split(',')
                for part in ligand_parts:
                    ligand_name, count_str = part.strip().split(':')
                    count = int(count_str.strip())
                    if count <= 0:
                        raise ValueError("数量必须是正整数")
                    
                    # Protenix 命令使用 'ligand'，其他命令使用 'ccd'
                    ligand_key = 'ligand' if isinstance(self, ProtenixCommand) else 'ccd'
                    ligand_data: Dict[str, Any] = {
                        ligand_key: ligand_name,
                        'count': count
                    }
                    if add_id:
                        ligand_data['id'] = ligand_name # Boltz 使用这个作为ID
                    ccd_ligands.append({'ligand': ligand_data})
            except ValueError as e:
                raise InputGeneratorError(
                    f"CCD配体参数格式错误或数量无效，请使用'配体CCD名称:个数'格式且数量为正整数: {str(e)}",
                    error_code=203
                ) from e
        return ccd_ligands

    def _get_first_sequence_stem(self) -> str:
        """
        从第一个提供的序列文件路径中提取文件名（不带扩展名）作为基础名称。

        Returns:
            str: 提取到的基础名称。

        Raises:
            InputGeneratorError: 当未提供任何序列文件时抛出。
        """
        sequence_args = [self.args.protein, self.args.dna, self.args.rna]
        for arg in sequence_args:
            if arg:
                # 只取第一个文件路径，例如 "file1.fasta:2,file2.fasta:3" -> "file1.fasta"
                first_item = arg.split(',')[0].strip()
                path = first_item.split(':')[0].strip()
                return Path(path).stem
        raise InputGeneratorError("未提供任何序列文件，无法确定基础名称。", error_code=204)


class BoltzCommand(BaseCommand):
    """
    生成Boltz-2 YAML输入文件的命令。
    """
    @staticmethod
    def add_arguments(parser: argparse.ArgumentParser) -> None:
        """
        为Boltz子命令添加特有和通用的命令行参数。
        """
        parser.add_argument(
            '--protein',
            help='蛋白质FASTA文件路径和数量，格式为"路径:数量,路径:数量" (e.g., "file1.fasta:2,file2.fasta:3")'
        )
        parser.add_argument(
            '--dna',
            help='DNA FASTA文件路径和数量，格式同上'
        )
        parser.add_argument(
            '--rna',
            help='RNA FASTA文件路径和数量，格式同上'
        )
        parser.add_argument(
            '--ccd',
            help='CCD配体，格式为"配体CCD名称:个数,配体CCD名称:个数" (e.g., "CCD_P4G:1,CCD_6OI:2")'
        )
        parser.add_argument(
            '--csv',
            required=True,
            help='包含化合物ID和SMILES的CSV文件（列：id, smiles）'
        )
        parser.add_argument(
            '--output-dir',
            default='.',
            help='YAML文件的输出目录（默认：当前目录）'
        )

    def run(self) -> None:
        """
        执行Boltz文件生成逻辑。
        """
        logger.info("开始处理Boltz输入文件生成...")
        
        # Boltz 需要 ID，所以传递 add_id=True
        sequences, _ = self._collect_sequences(add_id=True)
        ccd_ligands = self._parse_ccd_ligands(add_id=True)
        compounds = read_compound_csv(self.args.csv)
        base_name = self._get_first_sequence_stem()
        
        output_dir = Path(self.args.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        logger.info(f"输出目录: {output_dir.absolute()}")

        for compound in compounds:
            self._generate_yaml(base_name, sequences, ccd_ligands, compound, output_dir)

        logger.info(f"Boltz输入文件生成完成，共处理{len(compounds)}个化合物。")

    def _generate_yaml(self, base_name: str, sequences: List[Dict[str, Any]],
                       ccd_ligands: List[Dict[str, Any]], compound: Dict[str, str], output_dir: Path) -> None:
        """
        生成单个Boltz YAML文件。

        Args:
            base_name (str): 输出文件名的基础前缀。
            sequences (List[Dict[str, Any]]): 序列信息列表。
            ccd_ligands (List[Dict[str, Any]]): CCD配体信息列表。
            compound (Dict[str, str]): 当前化合物的ID和SMILES。
            output_dir (Path): 输出目录路径。
        """
        # 合并所有序列和配体
        all_entities = sequences.copy()
        all_entities.extend(ccd_ligands)
        # 添加化合物配体 (Boltz中化合物通常是binder，id为B)
        all_entities.append({
            'ligand': {
                'id': 'B',  # Boltz通常使用'B'作为binder的ID
                'smiles': compound['smiles'],
                'count': 1
            }
        })

        yaml_data = {
            'version': 1,
            'sequences': all_entities,
            'properties': [
                {
                    'affinity': {
                        'binder': 'B'
                    }
                }
            ]
        }

        output_filename = f"{base_name}_{compound['id']}.yaml"
        output_path = output_dir / output_filename

        try:
            with open(output_path, 'w', encoding='utf-8') as f_out:
                yaml.dump(yaml_data, f_out, sort_keys=False, allow_unicode=True)
            logger.info(f"已生成Boltz YAML文件: {output_path}")
        except Exception as e:
            raise InputGeneratorError(f"写入YAML文件失败: {output_path}, 错误: {str(e)}", error_code=301) from e


class ProtenixCommand(BaseCommand):
    """
    生成Protenix JSON输入文件的命令。
    """
    @staticmethod
    def add_arguments(parser: argparse.ArgumentParser) -> None:
        """
        为Protenix子命令添加特有和通用的命令行参数。
        """
        parser.add_argument(
            '--protein',
            help='蛋白质FASTA文件路径和数量，格式为"路径:数量,路径:数量" (e.g., "file1.fasta:2,file2.fasta:3")'
        )
        parser.add_argument(
            '--dna',
            help='DNA FASTA文件路径和数量，格式同上'
        )
        parser.add_argument(
            '--rna',
            help='RNA FASTA文件路径和数量，格式同上'
        )
        parser.add_argument(
            '--ccd',
            help='CCD配体，格式为"配体CCD名称:个数,配体CCD名称:个数" (e.g., "CCD_P4G:1,CCD_6OI:2")'
        )
        parser.add_argument(
            '--csv',
            required=True,
            help='包含化合物ID和SMILES的CSV文件（列：id, smiles）'
        )
        parser.add_argument(
            '--output',
            default='input.json',
            help='输出JSON文件名字 (默认: input.json)'
        )

    def run(self) -> None:
        """
        执行Protenix文件生成逻辑。
        """
        logger.info("开始处理Protenix输入文件生成...")
        
        # Protenix 不需要 ID，所以传递 add_id=False
        sequences, _ = self._collect_sequences(add_id=False)
        ccd_ligands = self._parse_ccd_ligands(add_id=False)
        compounds = read_compound_csv(self.args.csv)
        
        # Protenix的name字段通常基于蛋白质名称，这里沿用原逻辑
        protein_base_name = self._get_first_sequence_stem()
        
        entries = self._generate_entries(protein_base_name, sequences, ccd_ligands, compounds)
        self._write_json_output(entries)
        
        # 生成随机种子文件
        seed_path = Path(self.args.output).parent / 'seed.txt'
        generate_random_seed(str(seed_path))

        logger.info(f"Protenix输入文件生成完成: {self.args.output}")

    def _generate_entries(self, protein_base_name: str, sequences: List[Dict[str, Any]],
                         ccd_ligands: List[Dict[str, Any]], compounds: List[Dict[str, str]]) -> List[Dict[str, Any]]:
        """
        生成所有Protenix条目。

        Args:
            protein_base_name (str): 蛋白质文件的基础名称，用于构建条目名称。
            sequences (List[Dict[str, Any]]): 序列信息列表。
            ccd_ligands (List[Dict[str, Any]]): CCD配体信息列表。
            compounds (List[Dict[str, str]]): 化合物信息列表。

        Returns:
            List[Dict[str, Any]]: 包含所有Protenix条目的列表。
        """
        entries: List[Dict[str, Any]] = []
        for compound in compounds:
            entry_sequences = sequences.copy()
            entry_sequences.extend(ccd_ligands)
            
            # 化合物配体，Protenix使用'ligand'而不是'smiles'
            entry_sequences.append({
                'ligand': {
                    'ligand': compound['smiles'],
                    'count': 1
                }
            })

            entries.append({
                'sequences': entry_sequences,
                'name': f"{protein_base_name}_{compound['id']}"
            })
        return entries

    def _write_json_output(self, entries: List[Dict[str, Any]]) -> None:
        """
        写入JSON输出文件。

        Args:
            entries (List[Dict[str, Any]]): 要写入JSON文件的条目列表。

        Raises:
            InputGeneratorError: 当写入JSON文件失败时抛出。
        """
        output_path = Path(self.args.output)
        try:
            output_path.parent.mkdir(parents=True, exist_ok=True)
            with open(output_path, 'w', encoding='utf-8') as f_out:
                json.dump(entries, f_out, indent=4, ensure_ascii=False)
            logger.info(f"已生成Protenix JSON文件: {output_path}")
        except Exception as e:
            raise InputGeneratorError(f"写入JSON文件失败: {output_path}, 错误: {str(e)}", error_code=302) from e


def main() -> None:
    """
    主函数，解析命令行参数并执行相应命令。
    """
    parser = argparse.ArgumentParser(
        description='为Boltz-2或Protenix生成输入文件',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    subparsers = parser.add_subparsers(
        dest='command',
        required=True,
        help='选择要生成的输入文件类型'
    )

    # Boltz子命令
    boltz_parser = subparsers.add_parser(
        'boltz',
        help='生成Boltz-2 YAML输入文件',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    BoltzCommand.add_arguments(boltz_parser)
    boltz_parser.set_defaults(command_class=BoltzCommand)

    # Protenix子命令
    protenix_parser = subparsers.add_parser(
        'protenix',
        help='生成Protenix JSON输入文件',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    ProtenixCommand.add_arguments(protenix_parser)
    protenix_parser.set_defaults(command_class=ProtenixCommand)

    args = parser.parse_args()

    try:
        # 执行命令
        command = args.command_class(args)
        command.run()
    except InputGeneratorError as e:
        logger.error(f"操作失败: {e}")
        sys.exit(e.error_code) # 使用自定义错误码退出
    except Exception as e:
        logger.error(f"发生未知错误: {str(e)}")
        sys.exit(1)


if __name__ == '__main__':
    main()